import numpy as np
from scipy.special import softmax
import cv2 as cv
import io
import os
import time


class HairChangeColor(object):
    def __init__(self, model):
        self.net = model

    @classmethod
    def from_pretrained(cls, model_path:str="/models/.cache/hair_segmentation/hair_segmentation.tflite"):
        model = cv.dnn.readNet(model_path)
        return cls(model=model)

    def _mix_prev_mask(self, prev_mask, new_mask):
        combine_with_prev_ratio = 0.5
        eps = 1e-3
        uncertainty_alpha = 1.0 + (new_mask * np.log(new_mask + eps) + (
            1.0 - new_mask) * np.log(1.0 - new_mask + eps)) / np.log(2.0)
        uncertainty_alpha = np.clip(uncertainty_alpha, 0, 1)
        uncertainty_alpha *= 2.0 - uncertainty_alpha

        mixed_mask = new_mask * uncertainty_alpha + \
            prev_mask * (0.6 - uncertainty_alpha)
        return mixed_mask * combine_with_prev_ratio + (1.0 - combine_with_prev_ratio) * new_mask

    def strength_mask_color(self, mask, b_strength=1, g_strength=1, r_strength=1):
        coefficients = [b_strength, g_strength, r_strength]  # 每个通道的系数

        # 创建新的蒙版图，对应通道乘以不同的系数
        new_mask = np.copy(mask)
        for i in range(3):
            new_mask[..., i] = np.clip(
                mask[..., i] * coefficients[i], 0, 255).astype(mask.dtype)
        return new_mask

    def change_color(self, image_path, save_dir, color, b_strength=1, g_strength=1, r_strength=1, num_runs=2):
        if type(color) == str:
            # 去除颜色码中的'#'符号
            color = color.lstrip('#')
            # 将十六进制颜色码转换为RGB值
            color = tuple(int(color[i:i+2], 16) for i in (0, 2, 4))
            # 将RGB值映射到0-255范围 reverse -> cv2使用bgr
            color = list(int((value / 255) * 100) for value in color)
            color.reverse()
            color = tuple(color)

        frame = cv.imread(image_path)
        prev_mask = np.zeros((512, 512), dtype=np.float32)
        color = np.ones(frame.shape, dtype=np.uint8) * color
        # Prepare input
        blob = cv.dnn.blobFromImage(frame, 1.0 / 255, (512, 512), swapRB=True)
        blob = np.concatenate(
            (blob, prev_mask.reshape(1, 1, 512, 512)), axis=1)
        for i in range(num_runs):
            # Copy previous frame mask to a new tensor
            blob[0, 3] = prev_mask
            # Run network
            self.net.setInput(blob)
            out = self.net.forward()
            out = softmax(out, axis=1)
            mask = out[0, 1]
            prev_mask = self._mix_prev_mask(prev_mask, mask)
        mask = cv.resize(prev_mask, (frame.shape[1], frame.shape[0]))
        lum = cv.cvtColor(frame, cv.COLOR_BGR2GRAY) / 255
        mask *= lum
        mask = np.repeat(np.expand_dims(mask, axis=-1), 3, axis=-1)
        new_mask = np.copy(mask)
        new_mask[np.any(new_mask != [0, 0, 0], axis=-1)] = [255, 255, 255]
        # mask = cv.merge([mask[:, :, 0], np.ones_like(mask[:, :, 0]), np.zeros_like(mask[:, :, 0])])
        mask = self.strength_mask_color(
            mask=mask, b_strength=b_strength, g_strength=g_strength, r_strength=r_strength)
        result = (mask * (color.astype(np.float32) - frame) +
                  frame).astype(np.uint8)
        filename = os.path.join(
            save_dir, "hair-change-"+str(time.time())+".png")
        cv.imwrite(filename, result)
        return filename
